#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
异步判题服务 - 支持并发处理多个代码提交
"""

import asyncio
import subprocess
import tempfile
import os
import json
import time
import signal
import uuid
from typing import Dict, List, Tuple, Optional
from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor
import threading
import queue
import logging
from dataclasses import dataclass
from enum import Enum

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class JudgeStatus(Enum):
    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"

@dataclass
class JudgeTask:
    """判题任务"""
    task_id: str
    user_id: int
    problem_id: int
    code: str
    language: str
    test_cases: List[Dict]
    time_limit: int
    memory_limit: int
    status: JudgeStatus
    result: Optional[Dict] = None
    created_at: float = None
    started_at: Optional[float] = None
    completed_at: Optional[float] = None
    error_message: Optional[str] = None

class AsyncJudgeService:
    """异步代码评测服务"""
    
    def __init__(self, max_workers: int = 4, max_queue_size: int = 100):
        self.max_workers = max_workers
        self.max_queue_size = max_queue_size
        
        # 任务队列
        self.task_queue = queue.Queue(maxsize=max_queue_size)
        
        # 任务存储
        self.tasks: Dict[str, JudgeTask] = {}
        self.tasks_lock = threading.Lock()
        
        # 线程池
        self.executor = ThreadPoolExecutor(max_workers=max_workers)
        
        # 支持的编程语言配置
        self.supported_languages = {
            'python': {
                'extension': '.py',
                'command': 'python',
                'timeout': 10000,  # 10秒，以毫秒为单位
                'memory_limit': 64  # 默认64MB内存限制
            },
            'cpp': {
                'extension': '.cpp',
                'command': 'g++',
                'compile_command': ['g++', '-o', '{output}', '{source}'],
                'run_command': './{executable}',
                'timeout': 5000,  # 5秒，以毫秒为单位
                'memory_limit': 128
            },
            'java': {
                'extension': '.java',
                'command': 'javac',
                'compile_command': ['javac', '{source}'],
                'run_command': 'java {class_name}',
                'timeout': 5000,  # 5秒，以毫秒为单位
                'memory_limit': 128
            }
        }
        
        # 启动工作线程
        self.workers = []
        self.running = True
        self._start_workers()
        
        logger.info(f"异步判题服务已启动，最大工作线程数: {max_workers}")
    
    def _start_workers(self):
        """启动工作线程"""
        for i in range(self.max_workers):
            worker = threading.Thread(target=self._worker_loop, args=(i,))
            worker.daemon = True
            worker.start()
            self.workers.append(worker)
            logger.info(f"工作线程 {i} 已启动")
    
    def _worker_loop(self, worker_id: int):
        """工作线程主循环"""
        logger.info(f"工作线程 {worker_id} 开始运行")
        
        while self.running:
            try:
                # 从队列获取任务
                task = self.task_queue.get(timeout=1)
                if task is None:
                    continue
                
                logger.info(f"工作线程 {worker_id} 开始处理任务 {task.task_id}")
                
                # 更新任务状态
                with self.tasks_lock:
                    task.status = JudgeStatus.RUNNING
                    task.started_at = time.time()
                
                # 执行判题
                try:
                    result = self._judge_submission_sync(task)
                    with self.tasks_lock:
                        task.status = JudgeStatus.COMPLETED
                        task.result = result
                        task.completed_at = time.time()
                    
                    logger.info(f"工作线程 {worker_id} 完成任务 {task.task_id}")
                    
                except Exception as e:
                    logger.error(f"工作线程 {worker_id} 处理任务 {task.task_id} 失败: {str(e)}")
                    with self.tasks_lock:
                        task.status = JudgeStatus.FAILED
                        task.error_message = str(e)
                        task.completed_at = time.time()
                
                finally:
                    self.task_queue.task_done()
                    
            except queue.Empty:
                continue
            except Exception as e:
                logger.error(f"工作线程 {worker_id} 发生错误: {str(e)}")
                continue
        
        logger.info(f"工作线程 {worker_id} 已停止")
    
    def submit_judge_task(self, user_id: int, problem_id: int, code: str, 
                         language: str, test_cases: List[Dict], 
                         time_limit: int = None, memory_limit: int = None) -> str:
        """
        提交判题任务
        
        Returns:
            任务ID
        """
        # 生成任务ID
        task_id = str(uuid.uuid4())
        
        # 创建任务
        task = JudgeTask(
            task_id=task_id,
            user_id=user_id,
            problem_id=problem_id,
            code=code,
            language=language,
            test_cases=test_cases,
            time_limit=time_limit or self.supported_languages.get(language, {}).get('timeout', 5000),
            memory_limit=memory_limit or self.supported_languages.get(language, {}).get('memory_limit', 64),
            status=JudgeStatus.PENDING,
            created_at=time.time()
        )
        
        # 存储任务
        with self.tasks_lock:
            self.tasks[task_id] = task
        
        # 添加到队列
        try:
            self.task_queue.put(task, timeout=5)
            logger.info(f"任务 {task_id} 已提交到队列")
        except queue.Full:
            # 队列已满，删除任务
            with self.tasks_lock:
                del self.tasks[task_id]
            raise Exception("判题队列已满，请稍后重试")
        
        return task_id
    
    def get_task_status(self, task_id: str) -> Optional[Dict]:
        """获取任务状态"""
        with self.tasks_lock:
            task = self.tasks.get(task_id)
            if not task:
                return None
            
            return {
                'task_id': task.task_id,
                'status': task.status.value,
                'result': task.result,
                'error_message': task.error_message,
                'created_at': task.created_at,
                'started_at': task.started_at,
                'completed_at': task.completed_at
            }
    
    def get_queue_status(self) -> Dict:
        """获取队列状态"""
        return {
            'queue_size': self.task_queue.qsize(),
            'max_queue_size': self.max_queue_size,
            'active_workers': len([w for w in self.workers if w.is_alive()]),
            'max_workers': self.max_workers,
            'pending_tasks': len([t for t in self.tasks.values() if t.status == JudgeStatus.PENDING]),
            'running_tasks': len([t for t in self.tasks.values() if t.status == JudgeStatus.RUNNING]),
            'completed_tasks': len([t for t in self.tasks.values() if t.status == JudgeStatus.COMPLETED])
        }
    
    def _judge_submission_sync(self, task: JudgeTask) -> Dict:
        """同步判题（在工作线程中执行）"""
        try:
            if task.language not in self.supported_languages:
                return {
                    'status': 'CE',
                    'error_message': f'不支持的编程语言: {task.language}',
                    'score': 0,
                    'max_score': len(task.test_cases) * 10 if task.test_cases else 0,
                    'test_cases': []
                }
            
            lang_config = self.supported_languages[task.language]
            
            # 创建临时目录
            temp_dir = tempfile.mkdtemp()
            try:
                # 准备代码文件
                source_file = self._prepare_source_file(task.code, task.language, temp_dir)
                
                # 编译代码（如果需要）
                compile_result = self._compile_code(source_file, task.language, temp_dir)
                if compile_result['status'] != 'success':
                    return {
                        'status': 'CE',
                        'error_message': compile_result['message'],
                        'score': 0,
                        'max_score': len(task.test_cases) * 10 if task.test_cases else 0,
                        'test_cases': []
                    }
                
                # 运行测试用例
                test_results = []
                total_score = 0
                max_score = len(task.test_cases) * 10
                
                for i, test_case in enumerate(task.test_cases):
                    result = self._run_test_case(
                        source_file, task.language, test_case, 
                        task.time_limit, task.memory_limit, temp_dir, i
                    )
                    test_results.append(result)
                    
                    if result['status'] == 'AC':
                        total_score += 10
                
                # 确定最终状态
                if all(result['status'] == 'AC' for result in test_results):
                    final_status = 'AC'
                elif any(result['status'] == 'TLE' for result in test_results):
                    final_status = 'TLE'
                elif any(result['status'] == 'MLE' for result in test_results):
                    final_status = 'MLE'
                elif any(result['status'] == 'RE' for result in test_results):
                    final_status = 'RE'
                else:
                    final_status = 'WA'
                
                return {
                    'status': final_status,
                    'score': total_score,
                    'max_score': max_score,
                    'test_cases': test_results,
                    'execution_time': sum(result.get('execution_time', 0) for result in test_results),
                    'memory_usage': max(result.get('memory_usage', 0) for result in test_results)
                }
                
            finally:
                # 清理临时文件
                self._cleanup_temp_files(temp_dir)
                
        except Exception as e:
            logger.error(f"判题失败: {str(e)}")
            return {
                'status': 'RE',
                'error_message': str(e),
                'score': 0,
                'max_score': len(task.test_cases) * 10 if task.test_cases else 0,
                'test_cases': []
            }
    
    def _prepare_source_file(self, code: str, language: str, temp_dir: str) -> str:
        """准备源代码文件"""
        lang_config = self.supported_languages[language]
        extension = lang_config['extension']
        
        # 生成唯一文件名
        filename = f"code_{uuid.uuid4().hex[:8]}{extension}"
        filepath = os.path.join(temp_dir, filename)
        
        # 写入代码
        with open(filepath, 'w', encoding='utf-8') as f:
            f.write(code)
        
        return filepath
    
    def _compile_code(self, source_file: str, language: str, temp_dir: str) -> Dict:
        """编译代码"""
        if language == 'python':
            return {'status': 'success'}  # Python不需要编译
        
        lang_config = self.supported_languages[language]
        compile_command = lang_config.get('compile_command')
        
        if not compile_command:
            return {'status': 'success'}
        
        try:
            # 替换命令中的占位符
            if language == 'cpp':
                output_file = os.path.join(temp_dir, 'program')
                cmd = [arg.replace('{output}', output_file).replace('{source}', source_file) 
                       for arg in compile_command]
            elif language == 'java':
                cmd = [arg.replace('{source}', source_file) for arg in compile_command]
            else:
                cmd = compile_command
            
            # 执行编译
            result = subprocess.run(
                cmd,
                cwd=temp_dir,
                capture_output=True,
                text=True,
                timeout=30  # 编译超时30秒
            )
            
            if result.returncode == 0:
                return {'status': 'success'}
            else:
                return {
                    'status': 'error',
                    'message': f'编译错误:\n{result.stderr}'
                }
                
        except subprocess.TimeoutExpired:
            return {
                'status': 'error',
                'message': '编译超时'
            }
        except Exception as e:
            return {
                'status': 'error',
                'message': f'编译失败: {str(e)}'
            }
    
    def _run_test_case(self, source_file: str, language: str, test_case: Dict, 
                      time_limit: int, memory_limit: int, temp_dir: str, case_index: int) -> Dict:
        """运行单个测试用例"""
        try:
            input_data = test_case.get('input', '')
            expected_output = test_case.get('output', '')
            
            lang_config = self.supported_languages[language]
            run_command = lang_config.get('run_command', '')
            
            # 构建运行命令
            if language == 'python':
                cmd = ['python', source_file]
            elif language == 'cpp':
                executable = os.path.join(temp_dir, 'program')
                cmd = [executable]
            elif language == 'java':
                class_name = os.path.splitext(os.path.basename(source_file))[0]
                cmd = ['java', '-cp', temp_dir, class_name]
            else:
                cmd = [run_command]
            
            # 记录开始时间
            start_time = time.time()
            
            # 运行程序
            result = subprocess.run(
                cmd,
                input=input_data,
                capture_output=True,
                text=True,
                cwd=temp_dir,
                timeout=time_limit / 1000.0  # 转换为秒
            )
            
            # 计算执行时间
            execution_time = (time.time() - start_time) * 1000  # 转换为毫秒
            
            # 检查输出
            actual_output = result.stdout.strip()
            
            if result.returncode != 0:
                return {
                    'status': 'RE',
                    'input': input_data,
                    'expected_output': expected_output,
                    'actual_output': result.stderr,
                    'execution_time': execution_time,
                    'memory_usage': 0
                }
            
            if execution_time > time_limit:
                return {
                    'status': 'TLE',
                    'input': input_data,
                    'expected_output': expected_output,
                    'actual_output': actual_output,
                    'execution_time': execution_time,
                    'memory_usage': 0
                }
            
            if actual_output == expected_output:
                return {
                    'status': 'AC',
                    'input': input_data,
                    'expected_output': expected_output,
                    'actual_output': actual_output,
                    'execution_time': execution_time,
                    'memory_usage': 0
                }
            else:
                return {
                    'status': 'WA',
                    'input': input_data,
                    'expected_output': expected_output,
                    'actual_output': actual_output,
                    'execution_time': execution_time,
                    'memory_usage': 0
                }
                
        except subprocess.TimeoutExpired:
            return {
                'status': 'TLE',
                'input': test_case.get('input', ''),
                'expected_output': test_case.get('output', ''),
                'actual_output': '',
                'execution_time': time_limit,
                'memory_usage': 0
            }
        except Exception as e:
            return {
                'status': 'RE',
                'input': test_case.get('input', ''),
                'expected_output': test_case.get('output', ''),
                'actual_output': str(e),
                'execution_time': 0,
                'memory_usage': 0
            }
    
    def _cleanup_temp_files(self, temp_dir: str):
        """清理临时文件"""
        try:
            import shutil
            shutil.rmtree(temp_dir)
        except Exception as e:
            logger.warning(f"清理临时文件失败: {str(e)}")
    
    def shutdown(self):
        """关闭服务"""
        logger.info("正在关闭异步判题服务...")
        self.running = False
        
        # 等待所有工作线程完成
        for worker in self.workers:
            worker.join(timeout=5)
        
        # 关闭线程池
        self.executor.shutdown(wait=True)
        
        logger.info("异步判题服务已关闭")

# 全局判题服务实例
async_judge_service = AsyncJudgeService(max_workers=4, max_queue_size=100)
